#include "WalkerGroupCPU.h"

void WalkerGroupCPU::allocate(unsigned int numParticles, 
                              unsigned int numWalkers)
{
  mX[0] = new float[numWalkers * numParticles];
  mY[0] = new float[numWalkers * numParticles];
  mZ[0] = new float[numWalkers * numParticles];
  mX[1] = new float[numWalkers * numParticles];
  mY[1] = new float[numWalkers * numParticles];
  mZ[1] = new float[numWalkers * numParticles];
  mCurBuffer = new unsigned int[numWalkers];
  mRandomGen = new int*[numWalkers];
  mWaveFunction[0] = new float[numWalkers];
  mWaveFunction[1] = new float[numWalkers];
  mKinetic[0] = new float[numWalkers];
  mKinetic[1] = new float[numWalkers];
  mPotential[0] = new float[numWalkers];
  mPotential[1] = new float[numWalkers];
  mTotalKinetic = new double[numWalkers];
  mTotalPotential = new double[numWalkers];
}

WalkerGroupCPU::WalkerGroupCPU(WalkerGroup* initial, 
                               unsigned int numWalkers, 
                               unsigned int walkerStart)
{
  mNumParticles = initial->getNumParticles();
  mNumWalkers = numWalkers;
  
  this->allocate(mNumParticles, mNumWalkers);

  //Copy the initial configuration to each walker in this 
  for(unsigned int i = 0; i < numWalkers; i++)
  {
    initial->getX(&mX[0][i * mNumParticles]);
    initial->getY(&mY[0][i * mNumParticles]);
    initial->getZ(&mZ[0][i * mNumParticles]);

    mCurBuffer[i] = 0;
    mRandomGen[i] = init_sprng(SPRNG_LFG, 
                               walkerStart + i, 
                               gNumWalkers, 
                               gRandomSeed,
                               SPRNG_DEFAULT);
  }

  this->initialize();
}

WalkerGroupCPU::WalkerGroupCPU(unsigned int xDim,
                               unsigned int yDim,
                               unsigned int zDim)
{
  mNumParticles = xDim * yDim * zDim;
  mNumWalkers = 1;

  this->allocate(mNumParticles, mNumWalkers);

  mCurBuffer[0] = 0;
  mRandomGen[0] = init_sprng(SPRNG_LFG, 0, 1, gRandomSeed, SPRNG_DEFAULT);

  //Initialize a grid of particles
  unsigned int curParticle = 0;

  for(unsigned int i = 0; i < xDim; i++)
  {
    for(unsigned int j = 0; j < yDim; j++)
    {
      for(unsigned int k = 0; k < zDim; k++)
      {
        mX[0][curParticle] = i * GRID_SPACING;
        mY[0][curParticle] = j * GRID_SPACING;
        mZ[0][curParticle] = k * GRID_SPACING;
        curParticle++;
      }
    }
  }

  this->initialize();
}

WalkerGroupCPU::WalkerGroupCPU(const char* filename)
{
  FILE* file;
  
  file = fopen(filename, "r");

  if(file == NULL)
  {
    perror("Error: Could not open grid file.\n");
    exit(1);
  }

  if(fscanf(file, "%u %u", &mNumWalkers, &mNumParticles) != 2)
  {
    printf("Error: Invalid grid file format. Parse error line 1.\n");
    exit(1);
  }

  this->allocate(mNumParticles, mNumWalkers);

  for(unsigned int i = 0; i < mNumWalkers * mNumParticles; i++)
  {
    if(fscanf(file, 
              "%E %E %E\n", 
              &mX[0][i], 
              &mY[0][i], 
              &mZ[0][i]) != 3)
    {
      printf("Error: Grid file parse error line %d.\n", i + 2);
      exit(1);
    }
  }

  for(size_t i = 0; i < mNumWalkers; i++)
  {
    mCurBuffer[i] = 0;
  }

  fclose(file);
  
  this->initialize();

}

double WalkerGroupCPU::getEnergy()
{
  double totalEnergy = 0.0;

  for(unsigned int i = 0; i < mNumWalkers; i++)
  {
    totalEnergy += mTotalPotential[i];
    totalEnergy += mTotalKinetic[i];
  }

  totalEnergy /= (double)mNumIterations;

  return totalEnergy;
}

void WalkerGroupCPU::writeToDisk(const char* filename)
{
  FILE* file;

  file = fopen(filename, "w");

  if(file == NULL)
  {
    printf("Error %s:%d: Could not open file for writing.\n", 
           __FILE__, 
           __LINE__);
    exit(1);
  }

  if(fprintf(file, "%u %u\n", mNumWalkers, mNumParticles) < 0)
  {
    printf("Error %s:%d: Could not write to file.\n", __FILE__, __LINE__);
    exit(1);
  }

  for(size_t curWalker = 0; curWalker < mNumWalkers; curWalker++)
  {
    unsigned int curBuffer = mCurBuffer[curWalker];

    for(size_t curParticle = 0; curParticle < mNumParticles; curParticle++)
    {
      if(fprintf(file,
                 "%10E %10E %10E\n", 
                 mX[curBuffer][curParticle + curWalker * mNumParticles],
                 mY[curBuffer][curParticle + curWalker * mNumParticles],
                 mZ[curBuffer][curParticle + curWalker * mNumParticles]) < 0)
      {
        printf("Error %s:%d: Could not write to file.\n", __FILE__, __LINE__);
      }
    }
  }

  fclose(file);
}

void WalkerGroupCPU::getX(float* output)
{
  for(size_t i = 0; i < mNumParticles; i++)
  {
    output[i] = mX[mCurBuffer[0]][i];
  }
}

void WalkerGroupCPU::getY(float* output)
{
  for(size_t i = 0; i < mNumParticles; i++)
  {
    output[i] = mY[mCurBuffer[0]][i];
  }
}

void WalkerGroupCPU::getZ(float* output)
{
  for(size_t i = 0; i < mNumParticles; i++)
  {
    output[i] = mZ[mCurBuffer[0]][i];
  }
}

void WalkerGroupCPU::compute(cl_event* event)
{
  double p;

  this->perturb();

#pragma omp parallel for\
  schedule(static)
  for(unsigned int curWalker = 0; curWalker < mNumWalkers; curWalker++)
  {
    unsigned int currentBuffer = mCurBuffer[curWalker];
    unsigned int proposedBuffer = currentBuffer == 0 ? 1 : 0;

    GridCPUComputeTotalEnergy(0, 
                              &mX[proposedBuffer][curWalker * mNumParticles],
                              &mY[proposedBuffer][curWalker * mNumParticles],
                              &mZ[proposedBuffer][curWalker * mNumParticles],
                              mNumParticles,
                              &mWaveFunction[proposedBuffer][curWalker],
                              &mPotential[proposedBuffer][curWalker],
                              &mKinetic[proposedBuffer][curWalker]);
  
    p = exp(2.0 * (mWaveFunction[proposedBuffer][curWalker] - 
                   mWaveFunction[currentBuffer][curWalker]));

    float rand = sprng(mRandomGen[curWalker]);

    //If we accept, add the kinetic and potentials into the total and swap 
    //buffers.
    if(p >= 1.0 || rand < p)
    {
      mCurBuffer[curWalker] = proposedBuffer;

      mTotalKinetic[curWalker] += mKinetic[proposedBuffer][curWalker];
      mTotalPotential[curWalker] += mPotential[proposedBuffer][curWalker];
    }
    else
    {
      mTotalKinetic[curWalker] += mKinetic[currentBuffer][curWalker];
      mTotalPotential[curWalker] += mPotential[currentBuffer][curWalker];
    }
  }

  mNumIterations++;
}

void WalkerGroupCPU::perturb()
{
  for(unsigned int curWalker = 0; curWalker < mNumWalkers; curWalker++)
  {
    unsigned int currentBuffer = mCurBuffer[curWalker];
    unsigned int proposedBuffer = currentBuffer == 0 ? 1 : 0;

    for(unsigned int curParticle = 0; 
        curParticle < mNumParticles; 
        curParticle++)
    {
      mX[proposedBuffer][curWalker * mNumParticles + curParticle] = 
        mX[currentBuffer][curWalker * mNumParticles + curParticle] +
        (sprng(mRandomGen[curWalker]) - 0.5) *
        gStepSize;
      mY[proposedBuffer][curWalker * mNumParticles + curParticle] = 
        mY[currentBuffer][curWalker * mNumParticles + curParticle] +
        (sprng(mRandomGen[curWalker]) - 0.5) *
        gStepSize;
      mZ[proposedBuffer][curWalker * mNumParticles + curParticle] = 
        mZ[currentBuffer][curWalker * mNumParticles + curParticle] +
        (sprng(mRandomGen[curWalker]) - 0.5) *
        gStepSize;
    }
  }
}

void WalkerGroupCPU::initialize()
{
  static bool kernelInitialized = false;

  if(kernelInitialized == false)
  {
    try
    {
      GridCPUInit();
      kernelInitialized = true;
    }
    catch(GridException& err)
    {
      printf("%s\n", err.what());
      exit(1);
    }
  }

#pragma omp parallel for\
  schedule(static)
  for(unsigned int curWalker = 0; curWalker < mNumWalkers; curWalker++)
  {
    GridCPUComputeTotalEnergy(0, 
                              &mX[0][curWalker * mNumParticles],
                              &mY[0][curWalker * mNumParticles],
                              &mZ[0][curWalker * mNumParticles],
                              mNumParticles,
                              &mWaveFunction[0][curWalker],
                              &mPotential[0][curWalker],
                              &mKinetic[0][curWalker]);
    mCurBuffer[curWalker] = 0;
  }

}

WalkerGroupCPU::~WalkerGroupCPU()
{
  delete[] mX[0];
  delete[] mY[0];
  delete[] mZ[0];
  delete[] mX[1];
  delete[] mY[1];
  delete[] mZ[1];
  delete[] mCurBuffer;
  delete[] mRandomGen;
  delete[] mKinetic[0];
  delete[] mKinetic[1];
  delete[] mPotential[0];
  delete[] mPotential[1];
  delete[] mTotalKinetic;
  delete[] mTotalPotential;
  delete[] mWaveFunction[0];
  delete[] mWaveFunction[1];
}

